#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <bitset>
#include <gflags/gflags.h>
#include <glog/logging.h>

#include "core/grape.h"

using namespace std;

template<class VD, class ED, class MessageT, class ResultT, class TemporalT>
class BfsWorker: public Worker<VD, ED, MessageT, ResultT, TemporalT> {
	void pEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {

		int k = this->query.bound;

		VLOG(1) << "source = " << this->query.root << ", query-bound = " << k << endl;

		priority_queue<pair<int, VertexID>> queue;
		unordered_set<VertexID> result;

		if (fragment.isInnerVertex(this->query.root)) {
			queue.emplace(k, this->query.root);
		}

		while (!queue.empty()) {

			int distance = queue.top().first;
			VertexID u = queue.top().second;
			queue.pop();

			if (result.find(u) != result.end()) {
				continue;
			}

			result.insert(u);

			if (fragment.isBorderVertexWithOutgoingEdge(u)) {
				std::bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughOutgoingDest(u);
				for (int fid = 0; fid < dest.size(); fid++) {
					if (dest.test(fid)) {
						messageBuffer.add_message(fid, PairMessage<int>(u, distance));
					}
				}
			}

			if (distance > 0) {
				for (auto e : fragment.getOutgoingEdges(u)) {
					VertexID v = e.getDst();
					if (fragment.isInnerVertex(v)) {
						queue.emplace(distance - 1, v);
					}
				}
			}
		}

		partialResult = result;
	}

	void incEval(Fragment<VD, ED> &fragment, MessageBuffer<MessageT> &messageBuffer, ResultT &partialResult, TemporalT &temporal) {

		auto result = partialResult;
		priority_queue<pair<int, VertexID>> queue;

		unordered_set<VertexID> visited;

		for (auto item : messageBuffer.get_messages()) {
			VertexID u = item.getID();
			int distance = item.getValue();
			queue.emplace(distance, u);
		}

		while (!queue.empty()) {
			int distance = queue.top().first;
			VertexID u = queue.top().second;
			queue.pop();

			if (visited.find(u) != visited.end()) {
				continue;
			}

			visited.insert(u);
			if (fragment.isInnerVertex(u)) {
				result.insert(u);
			}

			if (fragment.isBorderVertexWithOutgoingEdge(u)) {
				std::bitset<NUM_FRAGMENTS> dest = fragment.getMsgThroughOutgoingDest(u);
				for (int fid = 0; fid < dest.size(); fid++) {
					if (dest.test(fid)) {
						messageBuffer.add_message(fid, PairMessage<int>(u, distance));
					}
				}
			}

			if (distance > 0) {
				for (auto e : fragment.getOutgoingEdges(u)) {
					VertexID v = e.getDst();
					if (fragment.isInnerVertex(v)) {
						queue.emplace(distance - 1, v);
					}
				}
			}
		}

		partialResult = result;
	}

	void assemble(vector<ResultT> &partialResults) {
		int count = 0;

		if (this->result_file != "") {
			std::ofstream fout;
			fout.open(this->result_file);

			for (auto result : partialResults) {
				for (auto item : result) {
					fout << item << endl;
				}
				count += result.size();
			}
			fout.close();
			VLOG(1) << " result.size  = " << count << endl;
		}
	}
};

int main(int argc, char *argv[]) {

	FLAGS_stderrthreshold = 0;

	google::SetUsageMessage("Usage: mpiexec.hydra [mpi_opts] ./bfs [grape_opts]");
	google::ParseCommandLineFlags(&argc, &argv, true);
	init_workers();

	WorkerParams param;

	param.algorithm = "bfs";
	param.enable_assemble = true;
	param.graph_load_strategy = ONLY_OUT;
	param.use_pattern = false;

	google::ShutDownCommandLineFlags();

	BfsWorker<char, char, PairMessage<int>, unordered_set<VertexID>, VertexID> bfsWorker;
	google::InitGoogleLogging("bfs");
	bfsWorker.run(param);
	google::ShutdownGoogleLogging();

	worker_finalize();
	return 0;
}
